#! py -3
# coding:utf-8

import cv2 as cv
import numpy as np
import random
import math
import matplotlib.pyplot as plt
from six.moves import xrange
import copy
import os
from os import path


# import torch

def blurImg(img: np.ndarray, masked: np.ndarray):
    s = img.shape
    m = 1.0
    edge = False
    for c in xrange(s[1]):
        m = 1.0
        edge = False
        for r in xrange(s[0]):
            if not edge:
                if img[r, c] > 80:
                    edge = True
            else:
                v = img[r, c]
                decay = random.random() * 0.025
                m = m * (1 - (0.005 + decay) * (v / 255.0))
                # m = m * (1 - 0.005* (v / 255.0))
                if m < 0.01:
                    img[r, c] = 0
                else:
                    if v * m > 15:
                        masked[r, c] = v
                    img[r, c] = v * m


def genCombinImg(img1: np.ndarray, img2: np.ndarray):
    scale = random.randint(30, 80)
    bgimg = np.random.rand(512, 512) * scale
    bgimg = bgimg.astype(np.uint8)
    lingimg = np.zeros_like(bgimg)

    if img1.size != img2.size:
        return bgimg, np.zeros_like(bgimg)

    if img1.shape != img2.shape:
        return bgimg, np.zeros_like(bgimg)

    n = random.randint(1, 8)
    for i in xrange(n):
        cols = lingimg.shape[1]
        j = random.randint(0, lingimg.shape[0] - 1)
        cnt = random.randint(1, 8)
        linea = np.random.randint(20, 100, (cnt,))
        # linea = np.random.normal(cnt / 2.0, 2.0, (cnt,))
        for v in linea:
            for k in xrange(cols):
                lingimg[j, k] = v * (0.6 + 0.4 * random.random())
            j = j + 1
            if j >= lingimg.shape[0]:
                break

    topcenter = [img1.shape[1] / 2, 0]
    p1 = [[1.0, 0], topcenter, [topcenter[0], 1.0]]
    p2 = copy.deepcopy(p1)
    scale = random.random() * 3
    p2[0][0] = -p2[0][0] * (scale + 1)
    scale = random.random() * 0.5
    p2[2][1] = p2[2][1] * (scale + 1)
    p1 = np.array(p1, np.float32)
    p2 = np.array(p2, np.float32)
    TM = cv.getAffineTransform(p1, p2)
    img1 = cv.warpAffine(img1, TM, (img1.shape[1], img1.shape[0]))
    img2 = cv.warpAffine(img2, TM, (img2.shape[1], img2.shape[0]))

    n = random.randint(0, img1.shape[1] - 512)
    img1 = img1[:, n:512 + n]
    img2 = img2[:, n:512 + n]

    n = random.randint(0, 3)
    if n == 1:
        img1 = cv.flip(img1, 0)
        img2 = cv.flip(img2, 0)
    elif n == 2:
        img1 = cv.flip(img1, 1)
        img2 = cv.flip(img2, 1)
    elif n == 3:
        img1 = cv.flip(img1, -1)
        img2 = cv.flip(img2, -1)

    n = random.randint(1, 10)
    if n > 9:
        center = (img1.shape[1] / 2, img1.shape[0] / 2)
        degree = random.randint(-5, 5)
        RM = cv.getRotationMatrix2D(center, degree, 1.0)
        img1 = cv.warpAffine(img1, RM, (img1.shape[1], img1.shape[0]))
        img2 = cv.warpAffine(img2, RM, (img2.shape[1], img2.shape[0]))

    top = random.randint(0, 512 - img1.shape[0])
    bottom = 512 - img1.shape[0] - top
    img1 = cv.copyMakeBorder(img1, top, bottom, 0, 0, cv.BORDER_CONSTANT, value=0)
    img2 = cv.copyMakeBorder(img2, top, bottom, 0, 0, cv.BORDER_CONSTANT, value=0)

    beta = random.random() * 0.5 + 0.5
    gamma = random.randint(0, 20)
    img1 = cv.addWeighted(bgimg, 1, img1, beta, gamma)
    img1 = cv.addWeighted(img1, 1, lingimg, 1, 0)
    return img1, img2


capture = cv.VideoCapture()
if capture.open("../fingertip.mp4") is not True:
    print("open video file failed!")
    exit()

savedir_input = "../oct_img/input_"
savedir_output = "../oct_img/output_"
filename_pre = "fingertip"
filename_i = 0

if not path.exists(savedir_input):
    os.makedirs(path.abspath(savedir_input))
if not path.exists(savedir_output):
    os.makedirs(path.abspath(savedir_output))

while (True):
    readok, img = capture.read()
    if (not readok) or (img is None):
        capture.set(cv.CAP_PROP_POS_FRAMES, 0)
        print("一次视频读取完毕")
        continue
        # break
    assert (type(img) is np.ndarray)
    img: np.ndarray
    img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    cv.imshow("org", img)
    key = cv.waitKey(1)
    # continue

    # img = img[100:-120, :]
    # img = 255 - img
    _, fil = cv.threshold(img, 30, 255, cv.THRESH_TOZERO)
    # edge = cv.GaussianBlur(fil, (3, 3), 1.5)
    # edge = cv.Canny(edge, 255, 0)
    fil: np.ndarray
    bimg = fil.copy()
    masked = np.zeros_like(bimg)
    blurImg(bimg, masked)

    # cv.imshow("edge", edge)
    cv.imshow("fil", fil)
    cv.imshow("masked", masked)
    cv.imshow("bimg", bimg)
    key = cv.waitKey(1)
    if key == 27:
        break

    # bimg, masked = genCombinImg(bimg, masked)
    bimg, fil = genCombinImg(bimg, fil)
    cv.imshow("filtered", fil)
    key = cv.waitKey(1)

    saveimg_mode = [cv.IMWRITE_PNG_COMPRESSION, 0, cv.IMWRITE_PNG_STRATEGY, cv.IMWRITE_PNG_STRATEGY_DEFAULT]
    filenamestr = savedir_input + "/" + filename_pre + "_%04d.png" % filename_i
    filenamestr = path.abspath(filenamestr)
    cv.imwrite(filenamestr, bimg, saveimg_mode)
    filenamestr = savedir_output + "/" + filename_pre + "_%04d.png" % filename_i
    filenamestr = path.abspath(filenamestr)
    # cv.imwrite(filenamestr, masked, saveimg_mode)
    cv.imwrite(filenamestr, fil, saveimg_mode)

    filename_i += 1

    cv.imshow("gen input", bimg)
    cv.imshow("gen output", masked)

    key = cv.waitKey(1)
    if key == 27:
        break

cv.destroyAllWindows()
